"""
LLM初始化数据脚本
创建默认的LLM模型配置
"""
import asyncio
import os
import sys
from sqlalchemy.orm import Session
from app.db.session import SessionLocal
from app.models.llm_model import LLMModel

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def create_default_models(db: Session):
    """创建默认的LLM模型配置"""
    
    # 检查是否已存在模型
    existing_count = db.query(LLMModel).count()
    if existing_count > 0:
        print("LLM模型已存在，跳过初始化")
        return
    
    # 默认模型配置
    default_models = [
        {
            "name": "GPT-3.5 Turbo",
            "provider": "openai",
            "model_name": "gpt-3.5-turbo",
            "api_key": os.getenv("OPENAI_API_KEY", ""),
            "base_url": None,
            "max_tokens": 4000,
            "temperature": 0.7,
            "is_active": bool(os.getenv("OPENAI_API_KEY"))
        },
        {
            "name": "GPT-4",
            "provider": "openai", 
            "model_name": "gpt-4",
            "api_key": os.getenv("OPENAI_API_KEY", ""),
            "base_url": None,
            "max_tokens": 4000,
            "temperature": 0.7,
            "is_active": bool(os.getenv("OPENAI_API_KEY"))
        },
        {
            "name": "Claude 3 Haiku",
            "provider": "anthropic",
            "model_name": "claude-3-haiku-20240307",
            "api_key": os.getenv("ANTHROPIC_API_KEY", ""),
            "base_url": None,
            "max_tokens": 4000,
            "temperature": 0.7,
            "is_active": bool(os.getenv("ANTHROPIC_API_KEY"))
        },
        {
            "name": "Claude 3 Sonnet",
            "provider": "anthropic",
            "model_name": "claude-3-sonnet-20240229",
            "api_key": os.getenv("ANTHROPIC_API_KEY", ""),
            "base_url": None,
            "max_tokens": 4000,
            "temperature": 0.7,
            "is_active": bool(os.getenv("ANTHROPIC_API_KEY"))
        },
        # {
        #     "name": "Gemini Pro",
        #     "provider": "google",
        #     "model_name": "gemini-pro",
        #     "api_key": os.getenv("GOOGLE_API_KEY", ""),
        #     "base_url": None,
        #     "max_tokens": 4000,
        #     "temperature": 0.7,
        #     "is_active": bool(os.getenv("GOOGLE_API_KEY"))
        # }
    ]
    
    for model_data in default_models:
        if model_data["is_active"] and model_data["api_key"]:
            model = LLMModel(**model_data)
            db.add(model)
            print(f"创建模型: {model_data['name']}")
    
    db.commit()
    print("LLM模型初始化完成")


def main():
    """主函数"""
    db = SessionLocal()
    try:
        create_default_models(db)
    except Exception as e:
        print(f"初始化失败: {str(e)}")
        db.rollback()
    finally:
        db.close()


if __name__ == "__main__":
    main()